source('code/estimation_functions.R')
library(tidyverse)
library(EnvStats)
library(intergraph)
library(statnet)
## Simulation steps:
## 1) Draw some orgs' ideal points
## 2) Draw some briefs' ideal points
## 3) Draw how many briefs each org signs
## 4) Assign orgs to briefs via weighted sampling

norgs = 1000
nbriefs = 1000
redblue = colorRampPalette(c("red","white","blue"))

generate_data = function(norgs, nbriefs){
    
  # Draw the orgs
  orgs = data.frame(id = paste0("A", 1:norgs),
                   ideal1 = EnvStats::rnormMix(norgs,mean1 = -0.75, sd1 = 0.25,
                                               mean2 = 0.75, sd2 = 0.25),
                   connections = ceiling(rexp(1000, 0.5))) # Here is the number 
                                                          # of briefs each org signs
  orgs$ideal2 = ideal2 = cut(orgs$ideal1, 10)
  
  # Draw the briefs
  briefs = data.frame(id = paste0("B", 1:nbriefs),
                    ideal1 = EnvStats::rnormMix(nbriefs,mean1 = -0.75, sd1 = 0.25,
                                                mean2 = 0.75, sd2 = 0.25))
  
  walk = data.frame(vals = levels(orgs$ideal2),
                    cols =   redblue(10))
  orgs = orgs %>% left_join(walk, by=c("ideal2" = "vals")) %>% select(-ideal2) 
  
  
  
  ## Produce an edgelist
  out = list()
  for(i in 1:nrow(orgs)){
    tmp = data.frame(a = orgs$id[i], b = sample(x = briefs$id, replace = FALSE,
                                       size = orgs$connections[i],
                                       prob = (1/abs(orgs$ideal1[i] - briefs$ideal1))^(1)
                                       ))
                                       ## Here is the weighted sampling by distance
    out[[i]] = tmp
  }
  out = do.call(rbind, out)
  out = rbind(out, data.frame(a = out$b, b = out$a))
  
  dat2 = as.network(x = out, multiple = TRUE)
  
  dat2%v%'ideal' = c(orgs$ideal1, rep(0, nbriefs))
  dat2%v%'color' = c(orgs$cols, rep("grey50", nbriefs))
  #plot(dat2, vertex.col = dat2%v%"color", usearrows=FALSE)
  return(dat2)
}

simple_sim = function(p){
  ## Make a sample amicus network
  g = generate_data(norgs, nbriefs)
  #plot(g, vertex.col = g%v%"color", usearrows=FALSE)
  
  truth_df = data.frame(name = network::network.vertex.names(g),
                        truth = g %v% 'ideal')
  
  ## Drop some nodes
  ## Drop ALL the briefs
  drops = c(rbinom(sum(grepl("A", truth_df$name)), 1, prob=p), rep(0, sum(grepl("B", truth_df$name))))
  g%v%'ideal2' = ifelse(drops, g%v%'ideal', NA) 
  g%v%'weight' = 1 
  
  truth_df = data.frame(name = network::network.vertex.names(g),
                        truth = g %v% 'ideal',
                        drops = as.logical(drops),
                        org = grepl("A", network::network.vertex.names(g)))
  
  ## Convert
  g2 = intergraph::asIgraph(g)
  g2 = g2 %>% set_edge_attr("weight", value=1)
  V(g2)$name = network::network.vertex.names(g)
  g3 = as_tbl_graph(g2)
  

  ## Do the estimation
  idealvec = g%v%'ideal'
  
  M = make_transition_matrix(g = g3,
                             exog = V(g3)$name[as.logical(drops)]) 
  test_scores = estimate_ideology_without_steady_state(M,
                                                       exog = V(g3)$name[as.logical(drops)],
                                                       score_vector = idealvec[as.logical(drops)])
  test_scores = test_scores[,1]
  summary(test_scores)
  
  result_df = data.frame(name = names(test_scores), pred = test_scores)
  result_df2 = merge(truth_df, result_df)
  
  out = result_df2 %>% filter(drops == 0, org==TRUE) %>%
    summarize(corr = cor(pred, truth),
              MAD = mad(pred, truth),
              samesign = mean(sign(pred)==sign(truth))) %>%
    mutate(dropprob =p)
  return(out)
}


set.seed(12345)
params = rep(seq(0.05, 0.95, by=0.05), each=20)

out = list()
for(i in 1:length(params)){
  out[[i]] = simple_sim(p = params[i])
  print(i)
}

out2 = do.call(rbind, out) %>% filter(dropprob<1)
save(out2, file="data/simple_sims.RData")